#include "blazepose_postprocess.h"

namespace sp::Model::BlazePose{

void registe_blazepose_postprocess(){
    blazepose_postprocess::registe();
}

blazepose_postprocess::blazepose_postprocess(std::vector<std::string> v): PythonFunction("/home/lx/SmartPipe/src/core/functions/Model/BlazePose/blazepose/", "/home/lx/SmartPipe/src/core/functions/Model/BlazePose/blazepose/repo/BlazePose_torch/", "blazepose", "blazepose", v){
    assert(v.size() == 0);
    name = "blazepose_postprocess";
    class_name = "Model-BlazePose-blazepose_postprocess";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({TENSOR},MESH));
}

blazepose_postprocess::blazepose_postprocess(): PythonFunction(){
    name = "blazepose_postprocess";
    class_name = "Model-BlazePose-blazepose_postprocess";
    pModulePath = "/home/lx/SmartPipe/src/core/functions/Model/BlazePose/blazepose/";
    pModuleHomePath = "/home/lx/SmartPipe/src/core/functions/Model/BlazePose/blazepose/repo/BlazePose_torch/";
    pModuleName = "blazepose";
    pClassName = "blazepose";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({TENSOR},MESH));
}

blazepose_postprocess::blazepose_postprocess(int width, int height): PythonFunction(){
    name = "blazepose_postprocess";
    class_name = "Model-BlazePose-blazepose_postprocess";
    pModulePath = "/home/lx/SmartPipe/src/core/functions/Model/BlazePose/blazepose/";
    pModuleHomePath = "/home/lx/SmartPipe/src/core/functions/Model/BlazePose/blazepose/repo/BlazePose_torch/";
    pModuleName = "blazepose";
    pClassName = "blazepose";
    this->width = width;
    this->height = height;
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({TENSOR},MESH));
}


blazepose_postprocess::~blazepose_postprocess(){

}

void blazepose_postprocess::registe(){
    Function::Register("Model-BlazePose-blazepose_postprocess", createObject);
    Function::Register("Model-BlazePose-blazepose_postprocess", createObject2);
}

Function* blazepose_postprocess::createObject(std::vector<std::string> params){
    return new blazepose_postprocess(params);
}

Function* blazepose_postprocess::createObject2(){
    return new blazepose_postprocess();
}

void blazepose_postprocess::start(){
    // 构建实例
    PythonFunction::defaultStart();
    args.push_back(std::to_string(width));
    args.push_back(std::to_string(height));
    assert(args.size() == 2);
    pArgs = PyTuple_New(args.size());
    for(size_t i = 0; i < args.size(); i++){
        PyTuple_SetItem(pArgs, i, Py_BuildValue("s", args[i].c_str()));
    }
    pInstance = PyObject_CallObject(pConstruct, pArgs);
}

bool blazepose_postprocess::waitForResource(){
    return true;
}

/*
    vector<vector<Tensor*>> -> [Tensor,Tensor,...]
        blazepose_postprocess: [Tensor,Tensor,...] -> [[String,String,...],[String,String,...],...]
    [[String,String,...],[String,String,...],...] -> vector<Head*> (Head* -> String* -> String* -> ...)
*/
void blazepose_postprocess::process(std::vector<std::vector<Data*>>& data_input, std::vector<Data*>& data_output){
    assert(data_input.size() != 0);
    assert(data_input[0].size() == 1);
    // 输入的tensor变成Python格式
    PyObject* data = convertToPython(data_input);
    // 调用Python函数
    pArgs = PyTuple_New(1);
    PyTuple_SetItem(pArgs, 0, data);
    pReturn = PyObject_CallMethod(pInstance, "postprocess", "O", pArgs);
    // 将结果解析成data并赋值给data_output
    convertToCpp(pReturn, data_input, data_output, 1);
    // 释放内存
    for(size_t i = 0; i < data_input.size(); i++){
        assert(data_input[i][0]->type == TENSOR);
        if(data_input[i][0]->subRefCount() == 0){
            executor->freeAll(data_input[i][0]);
        }
    }
}

bool blazepose_postprocess::releaseResource(){
    return true;
}

void blazepose_postprocess::finish(){
    PythonFunction::defaultFinish();
}

void blazepose_postprocess::copy(Function* other){
    Function::defaultCopy(other);
    blazepose_postprocess* oth = dynamic_cast<blazepose_postprocess*>(other);
    this->width = oth->width;
    this->height = oth->height;
}

}